package com.zy.config;

import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.*;

@Component
public class QdrantClient {

    private final WebClient client;
    private static final int VECTOR_DIM = 1024; // embedding 维度

    public QdrantClient(QdrantProperties props) {
        this.client = WebClient.builder()
                .baseUrl(props.getUrl())
                .defaultHeader("api-key", props.getApiKey()) // Qdrant API Key
                .build();
    }

    /** 插入 embedding */
    public void upsertVector(String collectionName, String id, List<Double> embedding, Map<String, Object> payload) {
        if (embedding.size() != VECTOR_DIM) {
            throw new IllegalArgumentException("embedding 长度必须为 " + VECTOR_DIM + "，当前长度=" + embedding.size());
        }

        // 转成 List<Float>
        List<Float> vectorList = new ArrayList<>(VECTOR_DIM);
        for (Double val : embedding) {
            vectorList.add(val.floatValue());
        }

        // 构建 point
        Map<String, Object> point = new HashMap<>();
        point.put("id", id != null ? id : UUID.randomUUID().toString());
        point.put("vector", vectorList);
        if (payload != null) {
            point.put("payload", payload);
        }

        Map<String, Object> body = Map.of("points", new Object[]{point});

        System.out.println("Qdrant 插入 body=" + body);

        client.put()
                .uri("/collections/" + collectionName + "/points")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToMono(String.class)
                .block();
    }
    /**
     * 搜索相似向量
     */
    public List<Map<String, Object>> searchSimilar(String collectionName, List<Double> embedding, int topK) {
        if (embedding.size() != VECTOR_DIM) {
            throw new IllegalArgumentException("embedding 长度必须为 " + VECTOR_DIM);
        }

        // ✅ 转 float[]
        float[] vector = new float[VECTOR_DIM];
        for (int i = 0; i < VECTOR_DIM; i++) {
            vector[i] = embedding.get(i).floatValue();
        }

        // ✅ 请求体
        Map<String, Object> body = Map.of(
                "vector", vector,
                "top", topK,
                "with_payload", true
        );

        // ✅ 发送请求
        Map<String, Object> response = client.post()
                .uri("/collections/" + collectionName + "/points/search")
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(body)
                .retrieve()
                .bodyToMono(Map.class)
                .block();

        // ✅ 返回结果
        return (List<Map<String, Object>>) response.get("result");
    }
}
